//
// Creater by mzy on 2022/5/6
// centralized 
//

#include <ros/ros.h>
#include "optitopo.h"
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/TwistStamped.h>
#include <std_msgs/Float32.h>
#include <string>
#include <ros/package.h>
#include <sensor_msgs/Imu.h>
#include <geometry_msgs/Vector3.h>
#include <geometry_msgs/Quaternion.h>
#include <geometry_msgs/QuaternionStamped.h>
#include <geometry_msgs/Point32.h>
#include <nav_msgs/Odometry.h>
#include <Eigen/Eigen>
#include <Eigen/Geometry>
#include <Eigen/Core>
#include <swarm_plan/UavsInfoStamped.h>
#include <swarm_plan/FormationInfo.h>
#include "gmm_map_new/GMMCommunication.h"

#include <visualization_msgs/MarkerArray.h>
#include <visualization_msgs/Marker.h>


using namespace Eigen;
using namespace std;

// global variable
double hover_height;
int Num_swarm = 4;
swarm_plan::UavsInfoStamped swarm_info_global;

int ctrl_rate = 30;
float servo_speed_scale = 2.0;// the scale map w speed of servo to servo command (by test)
ros::Publisher exp_pose_pub;
ros::Publisher exp_vel_pub;
ros::Publisher exp_acc_pub;
ros::Publisher control_acc_pub;
ros::Publisher control_z_pub;
ros::Publisher visual_pub;

ros::Publisher formation_pub;
ros::Publisher target_yaw_pub;
ros::Publisher visual_marker_pub;

//topo param
double pso_dt = 0.02; // 第i维度上的时间步长
double pso_wstart = 0.9;		// 第i维度上的起始权重系数
double pso_wend = 0.4;			// 第i维度上的终止权重系数
double pso_C1 = 1.49445;		// 第i维度上的加速度因子
double pso_C2 = 1.49445;		// 第i维度上的加速度因子
int pso_particle_num = 20;      // pso粒子个数
int pso_max_iter_num = 50;      // pso最大迭代次数

Matrix<double, 2, 4> swarm_init_pos;
Matrix<double, 2, 4> swarm_cur_pos;
Matrix<double, 2, 4> formation_param_bound;
Vector4d formation_param_step;
Vector4d formation_param_vel;
double formation_param_k_form = 0.1;

Vector2d center_swarm_init;
Vector3d center_swarm_cur;
Vector4d cam_yaw;

std::vector<Eigen::Vector2d> obs_means;  // 障碍物中心坐标,相对于集群队形中心 !!!必须是相对的
std::vector<double> pi_vector;           // 障碍物的点云数量/所占概率 TODO

bool flag_posinit = false;

double k_con = 1.0;
double k_fit = 1.0;
double k_env = 1.0;

double Plan_t = 1.0;

double fov_max = 1.570796;
double fov_min = -1.570796;
int fov_resolution = 90;
double env_gain_perish_time = 2.0;
double env_gain_perish_cof = 1.0;

Matrix<double, 2, 4> cam_yaw_bound;
double cam_yaw_step = 0.5235988;

geometry_msgs::PoseStamped env_potential_msg;
geometry_msgs::QuaternionStamped by_formation_msg0;
geometry_msgs::QuaternionStamped by_formation_msg1;
geometry_msgs::QuaternionStamped by_formation_param_msg;
geometry_msgs::QuaternionStamped Info_gainandexe_msg;

// tool func
template<typename T>
void readParam(ros::NodeHandle &nh, std::string param_name, T& loaded_param) {
    // template to read param from roslaunch
    const string& node_name = ros::this_node::getName();
    param_name = node_name + "/" + param_name;
    if (!nh.getParam(param_name, loaded_param)) {
        ROS_ERROR_STREAM("Failed to load " << param_name << ", use default value");
        //TODO:print default value
    }
    else{
        ROS_INFO_STREAM("Load " << param_name << " success");
        //TODO:print loaded value
    }
}


void loadRosParams(ros::NodeHandle &nh)
{
    readParam<double>(nh, "hover_height", hover_height);

    readParam<double>(nh, "k_con", k_con);
    readParam<double>(nh, "k_fit", k_fit);
    readParam<double>(nh, "k_env", k_env);

    readParam<double>(nh, "Plan_t", Plan_t);

    readParam<double>(nh, "fov_max", fov_max);
    readParam<double>(nh, "fov_min", fov_min);
    readParam<int>(nh, "fov_resolution", fov_resolution);
    readParam<double>(nh, "env_gain_perish_time", env_gain_perish_time);
    readParam<double>(nh, "env_gain_perish_cof", env_gain_perish_cof);

    readParam<double>(nh, "cam_yaw_step", cam_yaw_step);

    readParam<double>(nh, "formation_param_k_form", formation_param_k_form);

    readParam<double>(nh, "pso_dt", pso_dt);
    readParam<double>(nh, "pso_wstart", pso_wstart);
    readParam<double>(nh, "pso_wend", pso_wend);
    readParam<double>(nh, "pso_C1", pso_C1);
    readParam<double>(nh, "pso_C2", pso_C2);
    readParam<int>(nh, "pso_particle_num", pso_particle_num);
    readParam<int>(nh, "pso_max_iter_num", pso_max_iter_num);

    //read param from roslaunch
    vector<double> poses(8);
    vector<double> param_vec(8);
    vector<double> param_step_vec(4);
    vector<double> param_vel_vec(4);
    vector<double> cam_yaw_bound_vec(8);
    readParam<vector<double>>(nh, "swarm_init_pos", poses);
    readParam<vector<double>>(nh, "formation_param_bound", param_vec);
    readParam<vector<double>>(nh, "formation_param_step", param_step_vec);
    readParam<vector<double>>(nh, "formation_param_vel", param_vel_vec);
    readParam<vector<double>>(nh, "cam_yaw_bound", cam_yaw_bound_vec);
    // give to global variable
    for (int i = 0; i < 4; i++) {
        //对4个无人机依次赋值初始位置
        swarm_init_pos(0, i) = poses[2 * i];
        swarm_init_pos(1, i) = poses[2 * i + 1];
        center_swarm_init(0) += poses[2 * i];
        center_swarm_init(1) += poses[2 * i + 1];
        formation_param_bound(0, i) = param_vec[2 * i];
        formation_param_bound(1, i) = param_vec[2 * i + 1];
        formation_param_step(i) = param_step_vec[i];
        formation_param_vel(i) = param_vel_vec[i];
        cam_yaw_bound(0, i) = cam_yaw_bound_vec[2 * i];
        cam_yaw_bound(1, i) = cam_yaw_bound_vec[2 * i + 1];
    }
    center_swarm_init(0) = center_swarm_init(0) / 4;
    center_swarm_init(1) = center_swarm_init(1) / 4;
    ROS_INFO_STREAM("center_swarm_init x:" << center_swarm_init(0) << "  y:" << center_swarm_init(1));
}


void Visualize_Formation(Matrix<double,2,4> _formation,Vector3d _center_pos){
    visualization_msgs::Marker points,line_list;
    points.header.frame_id = line_list.header.frame_id = "map";
    points.header.stamp = line_list.header.stamp = ros::Time::now();
    points.ns = line_list.ns = "points_and_lines";
    points.action = line_list.action = visualization_msgs::Marker::ADD;
    points.pose.orientation.w = line_list.pose.orientation.w = 1.0;

    points.id = 100;
    line_list.id = 101;

    points.type = visualization_msgs::Marker::POINTS;
    line_list.type = visualization_msgs::Marker::LINE_LIST;

    points.scale.x = 0.15;
    points.scale.y = 0.15;

    line_list.scale.x = 0.07;
   
    points.color.a = 0.5;
    points.color.g = 0.0;
    points.color.r = 0.0;
    points.color.b = 0.0;

    line_list.color.a = 0.7;
    line_list.color.r = 0.4;
    line_list.color.g = 0.4;
    line_list.color.b = 0.6;

    for (int i = 0; i < 4;i++){
        geometry_msgs::Point p0;
        p0.x = _formation(0, i) + _center_pos(0);
        p0.y = _formation(1, i) + _center_pos(1);
        p0.z =  _center_pos(2);

        points.points.push_back(p0);
        for (int j = i; j < 4;j++){
            if (i == j){
                continue;
            }
            geometry_msgs::Point p1;
            p1.x = _formation(0, j) + _center_pos(0);
            p1.y = _formation(1, j) + _center_pos(1);
            p1.z =  _center_pos(2);
            line_list.points.push_back(p0);
            line_list.points.push_back(p1);
        }
    }
    visual_marker_pub.publish(points);
    visual_marker_pub.publish(line_list);
}


void PubCamYaw(Vector4d _target_yaw){
    geometry_msgs::Quaternion yaw_msgs;
    yaw_msgs.x = _target_yaw(0);
    yaw_msgs.y = _target_yaw(1);
    yaw_msgs.z = _target_yaw(2);
    yaw_msgs.w = _target_yaw(3);
    // yaw_msgs.x = 0.0;
    // yaw_msgs.y = 1.0;
    // yaw_msgs.z = 0.0;
    // yaw_msgs.w = -1.0;
    target_yaw_pub.publish(yaw_msgs); // TODO
}


void PubFormation(Matrix<double,2,4> _formation){
    swarm_plan::FormationInfo formation_msgs;
    formation_msgs.header.stamp = ros::Time::now();
    for (int i = 0; i < 4;i++){
        formation_msgs.poses[i].x = _formation(0, i);
        formation_msgs.poses[i].y = _formation(1, i);
        formation_msgs.poses[i].z = 0.0;
    }
    formation_pub.publish(formation_msgs);
}

void swarm_info_cb(const swarm_plan::UavsInfoStamped::ConstPtr& msg){
    swarm_info_global = *msg;
    center_swarm_cur = Vector3d::Zero();
    for (int i = 0; i < 4; i++)
    {
        center_swarm_cur(0) += swarm_info_global.poses[i].x;
        center_swarm_cur(1) += swarm_info_global.poses[i].y;
        center_swarm_cur(2) += swarm_info_global.poses[i].z;
        cam_yaw(i) = (double)swarm_info_global.PitchOfCamera[i].data;
    }
    center_swarm_cur(0) = center_swarm_cur(0) / 4;
    center_swarm_cur(1) = center_swarm_cur(1) / 4;
    center_swarm_cur(2) = center_swarm_cur(2) / 4;
    center_swarm_cur(0) += center_swarm_init(0);
    center_swarm_cur(1) += center_swarm_init(1);

    for (int i = 0; i < 4;i++){
        swarm_cur_pos(0, i) = swarm_info_global.poses[i].x - center_swarm_cur(0);
        swarm_cur_pos(1, i) = swarm_info_global.poses[i].y - center_swarm_cur(1);
    }
    if (!flag_posinit)
    {
        flag_posinit = true;
    }
}


void gmm_cb(const gmm_map_new::GMMCommunication& gmm_msg){
    if(!flag_posinit){
        return;
    }

    int obs_num = gmm_msg.gmm_number;
    if(obs_num == 0){
        return;
    }

    obs_means.clear();
    pi_vector.clear();
    // 筛选正前方的障碍物
    for (int i = 0; i < obs_num; i++)
    {
        // TODO:筛选标准
        if (center_swarm_cur(0) - 1 < gmm_msg.means[3 * i])
        {
            Vector2d tmp_obs;
            // 计算障碍相对于集群中心的位置
            tmp_obs(0) = gmm_msg.means[3 * i] - center_swarm_cur(0);
            tmp_obs(1) = gmm_msg.means[3 * i + 1] - center_swarm_cur(1);
            obs_means.push_back(tmp_obs);
            pi_vector.push_back((double)gmm_msg.pi[i]);
        }
    }
}


int main (int argc, char** argv) {
    // ros init
    ros::init(argc, argv, "swarm_plan", ros::init_options::AnonymousName);
    ros::NodeHandle nh;
    // load param
    loadRosParams(nh);

    //ros pub and sub
    //sub
    ros::Subscriber swarm_info_sub = nh.subscribe("/swarm_plan/swarm_info", 1, swarm_info_cb);
    ros::Subscriber gmm_sub = nh.subscribe("/gmm_mix_server", 1, gmm_cb);
    // pub
    target_yaw_pub = nh.advertise<geometry_msgs::Quaternion>("/swarm_plan/cam_yaw_target_angle", 1);
    formation_pub = nh.advertise<swarm_plan::FormationInfo>("/swarm_plan/formation", 1);
    exp_pose_pub = nh.advertise<geometry_msgs::PoseStamped>("/mavros/exp_pose/pose", 1);
    exp_vel_pub = nh.advertise<geometry_msgs::TwistStamped>("/mavros/exp_vel/vel", 1);
    exp_acc_pub = nh.advertise<geometry_msgs::TwistStamped>("/mavros/exp_acc/acc", 1);
    control_acc_pub = nh.advertise<geometry_msgs::TwistStamped>("/mavros/control_acc/acc", 1);
    control_z_pub = nh.advertise<geometry_msgs::TwistStamped>("/mavros/control_z/acc", 1);
    visual_pub = nh.advertise<nav_msgs::Odometry>("/visual_opti", 1);

    visual_marker_pub = nh.advertise<visualization_msgs::Marker>("/swarm_plan/visual_marker", 1);

    // 毕业论文
    ros::Publisher env_potential_pub=nh.advertise<geometry_msgs::PoseStamped>("/sawrm_plan/env_potential", 1);
    ros::Publisher by_formation0_pub=nh.advertise<geometry_msgs::QuaternionStamped>("/sawrm_plan/by_formation_01", 1);
    ros::Publisher by_formation1_pub=nh.advertise<geometry_msgs::QuaternionStamped>("/sawrm_plan/by_formation_23", 1);
    ros::Publisher by_formation_param_pub=nh.advertise<geometry_msgs::QuaternionStamped>("/sawrm_plan/by_formation_param", 1);
    ros::Publisher Info_gainandexe_pub=nh.advertise<geometry_msgs::QuaternionStamped>("/sawrm_plan/Info_gainandexe", 1);

    optitopo SwarmOptiTopo(Num_swarm, pso_particle_num, pso_max_iter_num,fov_resolution,fov_max,fov_min); //声明拓扑优化类
    SwarmOptiTopo.max_iter_num_ = pso_max_iter_num;	// 最大迭代次数
    SwarmOptiTopo.formation_init_ = swarm_init_pos;
    SwarmOptiTopo.formation_param_bound_ = formation_param_bound;
    SwarmOptiTopo.formation_param_step_ = formation_param_step;
    SwarmOptiTopo.formation_param_vel_ = formation_param_vel;
    SwarmOptiTopo.formation_param_k_form_ = formation_param_k_form;
    SwarmOptiTopo.k_con_ = k_con;
    SwarmOptiTopo.k_env_ = k_env;
    SwarmOptiTopo.k_fit_ = k_fit;
    SwarmOptiTopo.env_gain_perish_time_ = env_gain_perish_time;
    SwarmOptiTopo.env_gain_perish_cof_ = env_gain_perish_cof;
    SwarmOptiTopo.cam_yaw_bound_ = cam_yaw_bound;
    SwarmOptiTopo.cam_yaw_step_ = cam_yaw_step;
    for (int i = 0; i < SwarmOptiTopo.dim_; i++)
    {
        SwarmOptiTopo.dt_[i] = pso_dt; // 第i维度上的时间步长
        SwarmOptiTopo.wstart_[i] = pso_wstart;		// 第i维度上的起始权重系数
        SwarmOptiTopo.wend_[i] = pso_wend;			// 第i维度上的终止权重系数
        SwarmOptiTopo.C1_[i] = pso_C1;		// 第i维度上的加速度因子
        SwarmOptiTopo.C2_[i] = pso_C2;		// 第i维度上的加速度因子
    }

    ros::Rate ctrl_loop(ctrl_rate);
    ros::spinOnce();

    ros::Time pso_begin_time;
    ros::Time pso_begin_time_last = ros::Time::now();

    Vector4d param_out(1.0, 1.0, 0.0, 0.0);
    Vector4d param_test;
    Vector4d yaw_out;
    Matrix<double, 2, 4> formation_out;
    // main ctrl loop
    while(ros::ok()) {
        ros::spinOnce();
        ros::Time begin_time = ros::Time::now();

        SwarmOptiTopo.obs_means_.clear();
        SwarmOptiTopo.pi_vector_.clear();
        SwarmOptiTopo.obs_means_ = obs_means;
        SwarmOptiTopo.pi_vector_ = pi_vector;
        // for (int i = 0; i < pi_vector.size();i++){
        //     cout << "pi_vector: " << pi_vector[i]<< endl;
        // }

        // 是否利用现在位置进行闭环规划 TODO
        // SwarmOptiTopo.ConvertFormation2Param(swarm_init_pos, swarm_cur_pos, param_out);

        pso_begin_time = ros::Time::now();
        SwarmOptiTopo.GetOutput_pso(param_out, cam_yaw, (pso_begin_time-pso_begin_time_last).toSec());
        pso_begin_time_last = pso_begin_time;

        yaw_out = SwarmOptiTopo.alpha_out_;
        param_out = SwarmOptiTopo.formation_param_new_;
        formation_out = SwarmOptiTopo.formation_new_;
        PubCamYaw(yaw_out);
        PubFormation(formation_out);

        ros::Time end_time = ros::Time::now();
        ROS_INFO("\033[1;32m SwarmOptiTopo EXIT,  use time: %f ======== \033[0m",(end_time-begin_time).toSec());

        env_potential_msg.header.stamp = ros::Time::now();
        by_formation_msg0.header.stamp = ros::Time::now();
        by_formation_msg1.header.stamp = ros::Time::now();
        by_formation_param_msg.header.stamp = ros::Time::now();
        Info_gainandexe_msg.header.stamp = ros::Time::now();
        
        env_potential_msg.pose.position.x = SwarmOptiTopo.env_potential_total_;
        env_potential_msg.pose.position.y = SwarmOptiTopo.env_potential_form_;
        env_potential_msg.pose.position.z = SwarmOptiTopo.env_potential_total_ -  SwarmOptiTopo.env_potential_form_;
        by_formation_msg0.quaternion.w = formation_out(0, 0)+center_swarm_cur(0);
        by_formation_msg0.quaternion.x = formation_out(1, 0)+center_swarm_cur(1);
        by_formation_msg0.quaternion.y = formation_out(0, 1)+center_swarm_cur(0);
        by_formation_msg0.quaternion.z = formation_out(1, 1)+center_swarm_cur(1);
        by_formation_msg1.quaternion.w = formation_out(0, 2)+center_swarm_cur(0);
        by_formation_msg1.quaternion.x = formation_out(1, 2)+center_swarm_cur(1);
        by_formation_msg1.quaternion.y = formation_out(0, 3)+center_swarm_cur(0);
        by_formation_msg1.quaternion.z = formation_out(1, 3)+center_swarm_cur(1);
        by_formation_param_msg.quaternion.w = param_out(0);
        by_formation_param_msg.quaternion.x = param_out(1);
        by_formation_param_msg.quaternion.y = param_out(2);
        by_formation_param_msg.quaternion.z = param_out(3);
        Info_gainandexe_msg.quaternion.w = SwarmOptiTopo.Infoexe_;
        Info_gainandexe_msg.quaternion.x = SwarmOptiTopo.Infogain_;
        Info_gainandexe_msg.quaternion.y = SwarmOptiTopo.lambda2_;

        env_potential_pub.publish(env_potential_msg);
        by_formation0_pub.publish(by_formation_msg0);
        by_formation1_pub.publish(by_formation_msg1);
        by_formation_param_pub.publish(by_formation_param_msg);
        Info_gainandexe_pub.publish(Info_gainandexe_msg);

        // cout << "param_out" << param_out << endl;
        // cout << "center_swarm_cur" << center_swarm_cur << endl;
        // cout << "yaw_out" << yaw_out << endl;
        // cout << "param_out" << param_out << endl;
        // formation_out << 0.55, 0.4, -0.55, -0.4,
        //                 -0.3, 0.4, 0.3, -0.4;
        Visualize_Formation(formation_out, center_swarm_cur);

        ctrl_loop.sleep();
    }

    return 0;
}
